import asyncio
from contextvars import ContextVar, Token
from typing import Union
from sqlalchemy.pool import StaticPool
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    create_async_engine,
    async_scoped_session,
)
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.sql.expression import Update, Delete, Insert



session_context: ContextVar[str] = ContextVar("session_context")


def get_session_context() -> str:
    return session_context.get()


def set_session_context(session_id: str) -> Token:
    return session_context.set(session_id)


def reset_session_context(context: Token) -> None:
    session_context.reset(context)

# class RoutingSession(Session):
#     def get_bind(self, mapper=None, clause=None, **kw):
#         if self._flushing or isinstance(clause, (Update, Delete, Insert)):

#
#             return create_async_engine(configs.WRITER_DB_URL,
#                                        future=True,
#                                        echo=True,
#                                        pool_pre_ping=True,
#                                        pool_recycle=7200,
#                                        max_overflow=64,
#                                        # check_same_thread参数指定了SQLAlchemy引擎是否允许在不同线程之间共享数据库连接。
#                                        # 默认情况下，SQLAlchemy不允许在不同线程之间共享数据库连接，
#                                        # 因此设置check_same_thread=False参数可以允许在不同线程之间共享数据库连接。
#                                        connect_args={"check_same_thread": True},
#                                        # poolclass=StaticPool
#                                        )
#         else:
#             print("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa！！WRITER_DB_URL！！",
#                   configs.WRITER_DB_URL)
#             print("来sssssssssssssssssssssssssssssssssssssssss呗！！READER_DB_URL！！", configs.READER_DB_URL)
#             return create_async_engine(configs.READER_DB_URL,
#                                        future=True,
#                                        echo=True,
#                                        pool_pre_ping=True,
#                                        pool_recycle=7200,
#                                        max_overflow=64,
#                                        # check_same_thread参数指定了SQLAlchemy引擎是否允许在不同线程之间共享数据库连接。
#                                        # 默认情况下，SQLAlchemy不允许在不同线程之间共享数据库连接，
#                                        # 因此设置check_same_thread=False参数可以允许在不同线程之间共享数据库连接。
#                                        connect_args={"check_same_thread": True},
#                                        # poolclass=StaticPool
#                                        )
#
#
# core_app.add_middleware(SQLAlchemySessionContextMiddleware)